import akshare as ak
import pandas as pd
import numpy as np

# import talib as ta
from datetime import datetime


# ===================== 数据获取模块 =====================
def get_stock_data(stock_code, start_date, end_date):
    """获取个股日线数据（复权）及财务指标"""
    # 日线数据（含复权）
    price_df = ak.stock_zh_a_hist(
        symbol=stock_code,
        period="daily",
        start_date=start_date,
        end_date=end_date,
        adjust="hfq",
    )
    print("price_df=", price_df.head())
    price_df = price_df[["日期", "开盘", "最高", "最低", "收盘", "成交量"]].rename(
        columns={
            "日期": "date",
            "开盘": "open",
            "最高": "high",
            "最低": "low",
            "收盘": "close",
            "成交量": "volume",
        }
    )

    # 财务数据（ROE/毛利率）
    finance_df = ak.stock_financial_report_sina(stock_code)
    roe = finance_df[finance_df["指标"] == "净资产收益率"].iloc[0]["值"]  # 最新年报ROE
    gross_margin = finance_df[finance_df["指标"] == "销售毛利率"].iloc[0][
        "值"
    ]  # 毛利率

    return price_df, roe, gross_margin


def get_industry_data(industry_name):
    """获取行业成分股及市值数据"""
    industry_stocks = ak.stock_board_industry_cons_em(symbol=industry_name)
    print("industry_stocks=", industry_stocks.head())
    return industry_stocks[["代码", "名称", "成交额"]]


# ===================== 核心筛选逻辑 =====================
def is_volume_price_rising(data, window=5):
    """量价齐升：连续5日温和放量上涨[2](@ref)"""
    data = data.iloc[-window:]
    vol_increase = all(
        data["volume"] > data["volume"].shift(1).fillna(0)
    )  # 成交量连续递增
    price_increase = all(
        data["close"] > data["close"].shift(1).fillna(0)
    )  # 价格连续上涨
    vol_growth = all(
        1.01 < data["volume"] / data["volume"].shift(1) < 1.03
    )  # 温和放量(1-3%)
    return vol_increase and price_increase and vol_growth


def is_breakout(data):
    """平台突破：收盘价创30日新高+成交量放大至1.5倍均量[4](@ref)"""
    data["30d_high"] = data["high"].rolling(30).max()
    latest = data.iloc[-1]
    cond1 = latest["close"] > data["30d_high"].iloc[-2]  # 突破前高
    cond2 = latest["volume"] > 1.5 * data["volume"].rolling(5).mean().iloc[-2]  # 放量
    return cond1 and cond2


def calculate_dragon_score(stock_code, industry_df):
    """计算龙头股评分（行业地位+资金强度）[6,8](@ref)"""
    # 行业市值排名（前3为龙头）
    industry_df = industry_df.sort_values("总市值", ascending=False)
    rank = industry_df[industry_df["代码"] == stock_code].index[0] + 1
    rank_score = 1.0 if rank <= 3 else 0.5 if rank <= 5 else 0.2

    # 资金强度（大单净量）
    money_flow = ak.stock_individual_fund_flow(stock_code)
    big_net_ratio = money_flow.iloc[-1]["主力净占比"]  # 最新主力净流入占比
    capital_score = min(max(big_net_ratio / 20.0, 0), 1.0)  # 归一化到0-1

    return 0.6 * rank_score + 0.4 * capital_score  # 综合评分


# ===================== 主筛选函数 =====================
def select_dragon_stocks(industry_name, date):
    """筛选指定行业的龙头股"""
    # 步骤1：获取行业成分股
    industry_df = get_industry_data(industry_name)
    candidates = []

    for _, row in industry_df.iterrows():
        stock_code = row["代码"]
        # 步骤2：获取个股数据
        price_df, roe, gross_margin = get_stock_data(
            stock_code, date - pd.Timedelta(days=60), date
        )

        # 步骤3：财务过滤（ROE>15% + 毛利率>30%）[8](@ref)
        if roe < 15 or gross_margin < 30:
            continue

        # 步骤4：技术面筛选
        if not (is_volume_price_rising(price_df) and is_breakout(price_df)):
            continue

        # 步骤5：计算龙头评分
        dragon_score = calculate_dragon_score(stock_code, industry_df)

        # 步骤6：记录候选股
        candidates.append(
            {
                "代码": stock_code,
                "名称": row["名称"],
                "行业排名": industry_df[industry_df["代码"] == stock_code].index[0] + 1,
                "ROE": roe,
                "毛利率": gross_margin,
                "龙头评分": round(dragon_score, 2),
            }
        )

    # 按评分排序返回Top3
    return pd.DataFrame(candidates).sort_values("龙头评分", ascending=False).head(3)


# ===================== 执行示例 =====================
if __name__ == "__main__":
    # 配置参数
    target_industry = "光伏设备"  # 可替换为其他行业
    end_date = datetime.now().strftime("%Y%m%d")

    # 执行筛选
    dragon_stocks = select_dragon_stocks(target_industry, end_date)

    # 输出结果
    print(f"【{target_industry}行业龙头股Top3】")
    print(dragon_stocks[["代码", "名称", "行业排名", "龙头评分"]])
